Amazon SageMakerでLightGBMが使えるコンテナイメージを作ってみた

こんにちは、大阪DI部の大澤です。

Amazon SageMakerはAWSの機械学習のフルマネージドサービスです。ノートブックインスタンスによるデータ探索から機械学習モデルの学習から推論用エンドポイントへのモデルのデプロイなどを行うことができます。

SageMakerには幾つかの組み込みアルゴリズムがありますが、それ以外にもTensorFlowやMXNetなどのその他の機械学習フレームワークを使ったSageMaker上での学習やモデルのホスティング、バッチ変換なども可能です。いくつかのフレームワークについてはSageMakerで学習と推論を行うときに必要なコンテナイメージが用意されています。しかしながら、用意されているコンテナイメージでは要件が満たせない場合には、自らでコンテナイメージを作成するかマーケットプレイスでアルゴリズム/モデルパッケージをサブスクライブする必要があります。
今回は学習/推論用コンテナイメージを作成し、SageMakerでLightGBMを使ってみたいと思います。

学習/推論用コンテナイメージに関する基本的な解説は以下のエントリをご覧ください。

Amazon SageMakerで独自の学習/推論用コンテナイメージを作ってみる

やってみる

概要

• 学習/推論兼用のコンテナイメージを作成し、ECRのリポジトリへプッシュします。
  • 参考にするサンプルノートブック: amazon-sagemaker-examples/scikit_bring_your_own.ipynb at master · awslabs/amazon-sagemaker-examples
• LightGBMで分類モデルを学習させます。
  • データセット: Iris
• 学習させたもモデルを推論エンドポイントにホストさせて、推論リクエストを投げて確認します。

※SageMakerのノートブックインスタンス上で作業します。

学習の仕様

• データ(入力): lightgbm.Datasetのsave_binaryメソッドで書き出されるバイナリ形式のデータセット。
  • 対応するデータチャネル: train(必須)、validation
  • 対応する入力モード: Fileのみ
• ハイパーパラメータ(入力): ブースティングラウンド数や葉の数、学習率等といったパラメータ。学習時にhyperparametersに辞書形式(パラメータ名:値)で設定。
• モデル(出力): lightgbm.Boosterのsave_modelメソッドで書き出されるモデルデータ。
• CPUにのみ対応

推論の仕様

• モデル(入力): lightgbm.Boosterのsave_modelメソッドで書き出されるモデルデータ。
• データ(入力): CSV形式。行ごとに推論したい特徴量データがあり、1行でも複数行でもOK。各列の値は学習させたデータの形式と同じ。
• データ(出力): JSON形式。resultsの中に各クラスごとの確率値がリスト形式で入っている。行数は入力データと同じで、列数は分類対象のクラス数と同じ。
• CPUにのみ対応

学習/推論用コンテナイメージの作成

サンプルノートブックで紹介されている構成を基にしますが、幾つかのファイルの内容を変更します。

• Dockerfile: コンテナイメージを作成する際に参照する設定ファイル。LightGBMのセットアップ追加を始めとして全体的に変更。
• train: 学習用コンテナ起動時に実行されるファイル。学習用処理を全体的に変更。
• predictor.py: 推論処理の記述ファイル。推論用処理を全体的に変更。
• serve: 推論用コンテナ起動時に実行されるファイル。shebangの内容をpython3用に変更。

Dockerfile、train、predictor.pyについては大きく変更しているため、それぞれ内容を紹介します。

Dockerfile

コンテナイメージを作成する際に参照する設定ファイルです。

参考: amazon-sagemaker-examples/Dockerfile at master · awslabs/amazon-sagemaker-examples

FROM ubuntu:18.04

# 必要なパッケージをインストールします
# インストールによって生じる不要なファイルは削除します
RUN apt -y update && apt install -y --no-install-recommends \
    wget \
    python3-distutils \
    nginx \
    ca-certificates \
    libgomp1 \
    && apt clean

# pythonパッケージをインストールします
# キャッシュファイルは重たいので、削除しておきます
RUN wget https://bootstrap.pypa.io/get-pip.py && python3 get-pip.py && \
    pip install wheel numpy scipy scikit-learn pandas lightgbm flask gevent gunicorn && \
    rm -rf /root/.cache

# 環境変数の設定

# pythonからの出力をバッファしないようにする（ログが早く届くようにするため）
ENV PYTHONUNBUFFERED=TRUE
# .pycファイルを作成しないようにする
ENV PYTHONDONTWRITEBYTECODE=TRUE
# スクリプトがあるディレクトリへパスを通しておく(コンテナ起動時に実行するtrainとserveファイルのパスを意識しないで済むようにするため)
ENV PATH="/opt/program:${PATH}"

# スクリプトをコピーする
COPY lightgbm_container /opt/program

# 作業ディレクトリを設定する
WORKDIR /opt/program

train

学習用コンテナ起動時に実行される学習用処理が記述されたスクリプトです。

参考: amazon-sagemaker-examples/train at master · awslabs/amazon-sagemaker-examples

#!/usr/bin/env python3

import os
import json
import sys
import traceback
import lightgbm as lgb


# sagemakerがデータを渡すためにコンテナにマウントするパス
prefix = '/opt/ml/'
input_path = prefix + 'input/data'
output_path = os.path.join(prefix, 'output')
model_path = os.path.join(prefix, 'model')
param_path = os.path.join(prefix, 'input/config/hyperparameters.json')
inputdataconfig_path = os.path.join(prefix, 'input/config/inputdataconfig.json')


# 有効なデータチャネル(Fileモードのみ対応)
valid_channel_names = ['train', 'validation']


def train():
    print('Starting the training.')
    try:
        # ハイパーパラメータを読み込みます
        with open(param_path, 'r') as f:
            hyperparams = json.load(f)

        # 入力データコンフィグを読み込みます
        with open(inputdataconfig_path, 'r') as f:
            inputdataconfig = json.load(f)

        # 入力データを読み込みます。
        inputdata_dic = {}
        for channel_name in inputdataconfig.keys():
            assert channel_name in valid_channel_names, 'input data channel must be included in '+str(valid_channel_names)
            inputdata_path = os.path.join(input_path, channel_name, channel_name+'.bin')
            inputdata_dic[channel_name] = lgb.Dataset(inputdata_path)


        # light-gbmで学習
        model = lgb.train(
            hyperparams,
            inputdata_dic['train'],
            valid_sets= [inputdata_dic['validation']] if 'validation' in inputdata_dic else None
        )

        # モデルを保存
        model.save_model(os.path.join(model_path, 'lightgbm_model.txt'))
        print('Training complete.')

    except Exception as e:
        # 何かエラーが発生したら、その内容をfailureに吐き出すことで失敗理由を伝達する
        trc = traceback.format_exc()
        with open(os.path.join(output_path, 'failure'), 'w') as s:
            s.write('Exception during training: ' + str(e) + '\n' + trc)
        # 標準出力に出すことでログにも送る
        print('Exception during training: ' + str(e) + '\n' + trc, file=sys.stderr)
        # 0以外の値を返すことで実行失敗を伝える
        sys.exit(255)

if __name__ == '__main__':
    train()

    # 0を返すことで実行成功を伝える
    sys.exit(0)

predictor.py

推論時に呼び出される推論アプリケーション用処理が記述されたスクリプトファイルです。

参考: amazon-sagemaker-examples/predictor.py at master · awslabs/amazon-sagemaker-examples

import os
import json
import io
import flask

import numpy as np
import lightgbm as lgb

prefix = '/opt/ml/'
model_path = os.path.join(prefix, 'model', 'lightgbm_model.txt')

# モデルをラップするクラス
class ScoringService(object):
    model = None

    @classmethod
    def get_model(cls):
        """クラスが保持しているモデルを返します。モデルを読み込んでなければ読み込みます。"""
        if cls.model == None:
            cls.model = lgb.Booster(model_file=model_path)
        return cls.model

    @classmethod
    def predict(cls, input):
        """推論処理

        Args:
            input (array-like object): 推論を行う対象の特徴量データ"""
        clf = cls.get_model()
        return clf.predict(input)

# 推論処理を提供するflaskアプリとして定義
app = flask.Flask(__name__)

@app.route('/ping', methods=['GET'])
def ping():
    """ヘルスチェックリクエスト
    コンテナが正常に動いているかどうかを確認する。ここで200を返すことで正常に動作していることを伝える。
    """
    health = ScoringService.get_model() is not None

    status = 200 if health else 404
    return flask.Response(response='\n', status=status, mimetype='application/json')

@app.route('/invocations', methods=['POST'])
def transformation():
    """推論リクエスト
    CSVデータが送られてくるので、そのデータを推論する。推論結果をCSVデータに変換して返す。
    """
    data = None

    # CSVデータを読み込む
    if flask.request.content_type == 'text/csv':
        with io.StringIO(flask.request.data.decode('utf-8')) as f:
            data = np.loadtxt(f, delimiter=',')
    else:
        return flask.Response(response='This predictor only supports CSV data', status=415, mimetype='text/plain')

    print('Invoked with {} records'.format(data.shape[0]))

    # 推論実行
    predictions = ScoringService.predict(data)

    # jsonに変換し、レスポンスデータを作成
    result = json.dumps({'results':predictions.tolist()})

    # レスポンスを返す
    return flask.Response(response=result, status=200, mimetype='text/json')

変更したスクリプトの紹介は以上です。ここからノートブックインスタンス上でノートブックを作成し、作業を進めます。

コンテナイメージの作成とECRリポジトリへのプッシュ

以下のスクリプトを実行して学習/推論用コンテナイメージをビルドし、ECRのリポジトリへプッシュします。

%%sh

# アルゴリズム名
algorithm_name=sagemaker-lightgbm

# ファイルを実行可能にする
chmod +x lightgbm_container/train
chmod +x lightgbm_container/serve

# アカウントID取得
account=$(aws sts get-caller-identity --query Account --output text)

# リージョン名
region='ap-northeast-1'

# リポジトリarn
fullname="${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest"

# ECRのリポジトリが存在しなければ作成する
aws --region ${region} ecr describe-repositories --repository-names "${algorithm_name}" > /dev/null 2>&1

if [ $? -ne 0 ]
then
    aws --region ${region} ecr create-repository --repository-name "${algorithm_name}" > /dev/null
fi

# ECRへのログインコマンドを取得し、ログインする
$(aws ecr get-login --region ${region} --no-include-email)


# コンテナイメージをビルドする
docker build  -t ${algorithm_name} .
docker tag ${algorithm_name} ${fullname}

# ECRのリポジトリへプッシュする
docker push ${fullname}

学習と推論

先ほど作成したコンテナイメージを使ってIrisデータを分類するモデルの学習と推論を試してみます。

セットアップ

まずは作業を進めるにあたり必要となるlightgbmパッケージのインストールを行います。

!conda install -c conda-forge lightgbm -y

パッケージの読み込みとデータの保存先などの設定を行います。

import boto3
import re
import os
from os import path
import numpy as np
import pandas as pd
from sagemaker import get_execution_role
import sagemaker as sage
from sagemaker.predictor import csv_serializer
import lightgbm as lgb
from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn import metrics
import json


# 各データを保存するS3の場所
prefix = 'sagemaker/byom-lightgbm/'
bucket_name = 'bucket_name'

# 学習やエンドポイント作成時などに使用するIAMロール
role = get_execution_role()

# sagemaker用セッションの作成
sess = sage.Session()

データの準備

Irisデータセットを読み込み、学習用と検証用に分割します。

# irisデータを読み込む
iris = datasets.load_iris()

# 学習用と検証用にデータを分ける
train_x, validation_x, train_y, validation_y = train_test_split(iris.data, iris.target, test_size=0.2, stratify=iris.target)

LightGBMに最適化されたデータセットコンポーネントで保存します。

# lgb用データセットを作成する
train = lgb.Dataset(train_x, label=train_y)

# validationデータは学習用データと関連づける
validation = train.create_valid(validation_x, label=validation_y)

# ローカルの保存場所
train_data_local = './data/train.bin'
val_data_local = './data/validation.bin'

# バイナリ形式で保存する
train.save_binary(train_data_local)
validation.save_binary(val_data_local)

S3へアップロードします。

train_data_s3 = sess.upload_data(train_data_local, key_prefix=path.join(prefix, 'input/train'), bucket=bucket_name)
val_data_s3 = sess.upload_data(val_data_local, key_prefix=path.join(prefix, 'input/validation'), bucket=bucket_name)

学習

モデルの学習を行います。まずはハイパーパラメータとメトリクスの定義を行います。 今回は最低限のパラメータのみ設定していますが、使えるパラメータはLightGBMのドキュメントをご覧ください。

• Parameters — LightGBM documentation

メトリクスによって学習状態を確認できるようにするには、学習時に出力されるログに含まれるメトリクス値に対応する正規表現を定義する必要があります。

• Define Objective Metrics - Amazon SageMaker

# ハイパーパラメータ
params = dict(
    num_round = 10,
    objective = 'multiclass',
    num_class = len(iris.target_names)
)
# メトリクス
metric_definitions = [dict(
    Name = 'multilogloss',
    Regex = '.*\\[[0-9]+\\].*valid_[0-9]+\'s\\smulti_logloss: (\\S+)'
)]

インスタンスタイプやコンテナイメージなどの設定を行い、学習ジョブを開始します。

account = sess.boto_session.client('sts').get_caller_identity()['Account']
region = sess.boto_session.region_name

modelartifact_path = "s3://"+path.join(bucket_name, prefix, 'output')
model = sage.estimator.Estimator(
    '{}.dkr.ecr.{}.amazonaws.com/sagemaker-lightgbm:latest'.format(account, region), # コンテナイメージのarn
    role, # 使用するIAMロール
    1, # インスタンス数
    'ml.c4.2xlarge', # インスタンスタイプ
    output_path=modelartifact_path, # モデルの保存場所
    sagemaker_session=sess, # SageMakerのセッション
    metric_definitions=metric_definitions # メトリクスの定義
)

# ハイパーパラメータを設定
model.set_hyperparameters(**params)

# 入力データを設定し、学習ジョブを実行
model.fit(dict(
    train = train_data_s3,
    validation = val_data_s3
))

数分すると完了します。インスタンスの準備などの時間が大半で、学習処理自体は一瞬で終わります。

マネジメントコンソールから学習ジョブの詳細ページのMonitorの箇所を見ると、定義したメトリクスのグラフを見ることができます。 今回は学習処理が一瞬なので、メトリクスの変化は見れず最後の値のみしか表示されませんでした。

推論

先ほど学習させた分類モデルで推論を試してみます。 まずはエンドポイントを作成し、モデルをデプロイします。

predictor = model.deploy(1, 'ml.m4.xlarge', serializer=csv_serializer)

デプロイが完了したら、検証用データを投げて推論結果を受け取ります。

result = predictor.predict(validation_x)
result = json.loads(result)
result

推論結果を受け取れました。

混同行列で結果を確認してみます。

cm = metrics.confusion_matrix(validation_y, np.argmax(result['results'], axis=1))
cm

最後にエンドポイントを削除します。

sess.delete_endpoint(predictor.endpoint)

さいごに

Amazon SageMakerでLightGBMが使えるようにする学習/推論用のコンテナを作ってみました。入力パラメータの検査などを行わないなど、かなり雑な作りではありますが、それなりに簡単に作れました。また、自分で学習/推論用コンテナイメージを作って動かしてみることでSageMakerをより深く理解できました。ぜひ、皆さんも試してみてください〜。

お読みいただきありがとうございました〜!